from misc_defs import *

class CellTracker():
	def __init__(self, cell_type, cell_index, train=False, analyze=False, display=False, training_size=100):		
		self.cell_type = cell_type
		self.cell_index = cell_index
		self.training = train
		self.analyzing = analyze
		self.displaying = display
		
		#initial setup
		self.load_cell_data()
		if self.displaying:
			init_windows()
		
		#Run algorithms
		if self.training:
			self.train(training_size, n_train_frames)
		if self.analyzing:
			self.analyze()
		
	def load_cell_data(self):
		print "Loading cell data..."
		self.cell_frames, self.cell_params, self.training_params, self.blank_img = get_cellframes(self.cell_type, self.cell_index)
	def empty_img(self):
		return np.zeros(np.shape(self.blank_img))
	def run_detection(self, cell_frames, cell_params):
		brightness = cell_params['brightness']
		contrast = cell_params['contrast']
		canny_thresh1 = cell_params['canny_thresh1']
		canny_thresh2 = cell_params['canny_thresh2']
		kern_size = cell_params['kern_size']

		last_contour = None
		selected_contours = []
		deltas = []

		for img_orig in cell_frames:
			img_orig = adjust_image(img_orig,contrast=contrast,brightness=brightness)
			img_edges = cv2.Canny(img_orig.copy(), canny_thresh1, canny_thresh2) #find edges
			img_dilation = dilate(img_edges, kern_size) #dilate
			contours, hier = cv2.findContours(img_dilation.copy(), CONTOUR_TYPE, CONTOUR_APPROX) #find contours
			contour_idx, contour = choose_contour(contours,last_contour,img_frame(img_orig), method = 'biggest', use_rects=False)
			selected_contours.append(contour)
			delta = delta_contour(contour,last_contour)
			if delta > 0:	deltas.append(delta)
			last_contour = contour
			
			if self.displaying:
				self.display(img_orig,img_edges,img_dilation,contours,contour_idx,contour)
		
		return (selected_contours, deltas)
	
	def display(self,img_orig,img_edges,img_dilation,contours,contour_idx,contour):
		img_contours = self.empty_img()
		cv2.drawContours( img_contours, contours, ALL, CD_COLOR, thickness=CD_THICKNESS)
		img_chosen_contour = self.empty_img()
		cv2.drawContours( img_chosen_contour, contours, contour_idx, CD_COLOR, thickness=CD_THICKNESS)
		display_all([[img_orig],[img_edges],[img_dilation],[img_contours],[img_chosen_contour],[img_orig, contour]])
		pause_video()		
	def train(self, training_size, n_train_frames):
		avg_deltas = []
		params_all = []
		cell_frames_orig = self.cell_frames
		for tt in range(training_size):
			print "%i/%i"%(tt+1,training_size)
			cell_params = random_params(self.training_params, self.params_all)
			frame1 = random.randrange(n_train_frames,len(cell_frames_orig)) - n_train_frames
			cell_frames = cell_frames_orig[frame1:frame1+n_train_frames] #to analyze just subsection of images for speed
			selected_contours,deltas = self.run_detection(cell_frames, cell_params)
			if 1: #valid_contours(selected_contours,cell_frames):
				avg_deltas.append(np.average(deltas))
				params_all.append(cell_params)
		#Conclude training:
		if len(avg_deltas) == 0:
			print "No good parameters found. Try more training iterations or close in tighter on cell."
		else:
			best_params = params_all[avg_deltas.index(min(avg_deltas))]
			store_params(self.cell_type,self.cell_index,best_params)
	def analyze(self):
		selected_contours,deltas = self.run_detection(self.cell_frames,self.cell_params)
		
		#Length ratio
		cv2.namedWindow('debugging')
		lens = []
		for c in range(1,len(selected_contours)):
			print "%i / %i"%(c, len(selected_contours)-1)
			lens.append(length_ratio(selected_contours[c],self.empty_img()))
		lens = [l for l in lens if abs(l-np.average(lens))<3*np.std(lens)]
		print "Ratio: %0.3f" %(np.max(lens)/np.min(lens)) #np.max/min(lens) represents the cell's elongation value at most/least elongated moment	